import PySimpleGUI as Sg
import subprocess


# 创建索引类，传入索引、名称、描述、大小等信息
class WIMIndex:
    def __init__(self, index, name, description, size):
        self.index = index
        self.name = name
        self.description = description
        self.size = size


def format_size(size_str):
    """
    格式化文件大小，将字节转换为["B", "kB", "MB", "GB", "TB", "PB", "EB", "ZB", "YB"]
    param size: 文件大小的字节数
    return: 格式化后的文件大小
    """
    size_str = size_str.replace(',', '').replace('个字节', '')
    size = float(size_str)
    units = ["B", "kB", "MB", "GB", "TB", "PB", "EB", "ZB", "YB"]
    for i, unit in enumerate(units[:-1]):
        if size < 1024.0:
            return f"{size:.2f} {unit}"
        size /= 1024.0
    return f"{size:.2f} {units[-1]}"


def get_wim_indexes(wim_f):
    # 获取WIM文件中的索引信息，并返回一个包含索引信息的列表
    command = f'dism /Get-WimInfo /WimFile:{wim_f}'
    output = subprocess.check_output(command, shell=True, text=True)
    out_list = output.split("\n\n")
    wim_index_list = []
    for x in out_list:
        if "索引" in x:
            index = 0
            name = ""
            description = ""
            size = ""
            for row in x.splitlines():
                if "索引" in row:
                    index = row.split(":")[1].strip()
                if "名称" in row:
                    name = row.split(":")[1].strip()
                if "描述" in row:
                    description = row.split(":")[1].strip()
                if "大小" in row:
                    size = row.split(":")[1].strip()
            wim_index_list.append(WIMIndex(index, name, description, size))
            print(f"{index} {name} {format_size(size)}")
    return wim_index_list


layout = [
    [Sg.Multiline(key='-Index_Info-', size=(65, 10))],
    [Sg.Input(key='-WIM_FILE-', size=(53, 1)),
     Sg.FileBrowse('选择映像文件')],
    [Sg.Button('查询索引', key='-QUERY_INDEXES-'), Sg.Combo([], size=(55, 1), enable_events=True, key='-INDEX_COMBO-')],
    [Sg.Button('获取索引值', key='-Get_Index-'), Sg.InputText(key='-INDEX_VALUE-', size=(55, 1))],
]

window = Sg.Window('获取wim映像中的索引信息').Layout(layout)

while True:
    event, values = window.Read()

    # 用户点击“查询索引”按钮
    if event == '-QUERY_INDEXES-':
        wim_file = values['-WIM_FILE-']
        if wim_file:  # 确保WIM文件已选择
            indexes = get_wim_indexes(wim_file)
            # 将索引列表中每个索引对象的name和index填充到Combo中，格式为 "索引名 (索引号)"
            choices = [f"{index.index} {index.name} {format_size(index.size)}" for index in indexes]
            window['-Index_Info-'].Update(
                f"共找到{len(indexes)}个索引\n{'\n'.join(choices)}",
            )
            window['-INDEX_COMBO-'].Update(values=choices)
            # 将'-INDEX_COMBO-'的已选择默认值设为第一个
            window['-INDEX_COMBO-'].Update(value=choices[0])
        else:
            Sg.Popup('请先选择WIM文件')

    # 用户点击“获取索引值”按钮
    if event == '-Get_Index-':
        selected_index_combo_value = window['-INDEX_COMBO-'].Get()
        print(selected_index_combo_value)
        if selected_index_combo_value:
            # 分离出选中的索引号
            index_info = selected_index_combo_value.split()
            try:
                index_number = int(index_info[0].strip())
                window['-INDEX_VALUE-'].Update(index_number)
            except (IndexError, ValueError):
                Sg.Popup('无法从选择的索引中获取有效索引号')
        else:
            Sg.Popup('请先从下拉列表中选择一个索引')



    elif event in (None, 'Exit'):  # 如果用户关闭窗口，则退出程序
        break

window.Close()
